import time
from glob import glob
from multiprocessing import Process, Manager
from multiprocessing import freeze_support

import psutil
from tqdm import tqdm
from config import config
from check_data import gen_hash


# 等待方法
def wait_process():
    available_memory_percent = 100 - psutil.virtual_memory().percent
    if available_memory_percent < 10:
        while True:
            time.sleep(10)
            if available_memory_percent > 20:
                break


def process_data(signal_dict, data_list):
    batch_size = signal_dict.get("batch_size")

    for one in tqdm(range(0, len(data_list), batch_size)):
        # 检测等待
        wait_process()
        # 处理数据
        two = one + batch_size
        # 数据读取或者对齐
        get_data(data_list[one:two])


# 数据处理
def get_data(data_path):
    for one in data_path:
        hash_name = gen_hash(one)
        seq_len = config["seq_len"]
        with open(one, "r", encoding="utf-8") as f:
            one_data = f.read()
            one_data = "".join(one_data.split())
        for i_id, i in enumerate(range(0, len(one_data), seq_len)):
            j = seq_len + i
            two_data = one_data[i:j]
            if seq_len != len(two_data):
                two_data += "。" * (seq_len - len(two_data))
            with open(config["train_data_set_path"].replace("*", "") + "{}_{}.txt".format(hash_name, i_id), "w",
                      encoding="utf-8") as f:
                f.write(two_data)
    return one_data


# 控制程序
def main_data():
    path_list = glob(config["data_set_path"])

    num_works = 10
    batch_size = 128

    # 参数设置
    signal_dict = Manager().dict()

    signal_dict["batch_size"] = batch_size

    # 开启异步处理数据
    p_list = []
    for i in range(0, len(path_list), len(path_list) // num_works):
        j = i + len(path_list) // num_works
        p = Process(target=process_data, args=(signal_dict, path_list[i:j]))
        p.start()
        p_list.append(p)
    # 训练结束回收进程
    for p in p_list:
        p.join()


if __name__ == '__main__':
    freeze_support()

    main_data()
